from stat import S_ISDIR
from paramiko import SFTPClient, Transport, RSAKey, DSSKey, ECDSAKey, Ed25519Key


class SFTP(object):
    def _get_pkey(self, pkey_obj, ssh_key_content, password=None):
        return pkey_obj.from_private_key(ssh_key_content, password=password)

    def connect(self, host, username, password=None, ssh_key_content=None, port=22, timeout=30):
        transport = Transport(sock=(host, port))
        transport._active_check_timeout = timeout

        if ssh_key_content:
            pkey = self._get_pkey(RSAKey, ssh_key_content, password=password) or \
                   self._get_pkey(DSSKey, ssh_key_content, password=password) or \
                   self._get_pkey(ECDSAKey, ssh_key_content, password=password) or \
                   self._get_pkey(Ed25519Key, ssh_key_content, password=password)
            transport.connect(pkey=pkey)
        else:
            transport.connect(username=username, password=password)

        self._sftp_obj = SFTPClient.from_transport(transport)

    def get_dir_items(self, parent_dir):
        dir_items = []
        for dir_item in self._sftp_obj.listdir_attr(parent_dir):
            if not dir_item.filename.startswith("."):
                dir_items.append({
                    "parent": parent_dir,
                    "name": dir_item.filename,
                    "size": dir_item.st_size,
                    "mtime": dir_item.st_mtime,
                    "type": "d" if S_ISDIR(dir_item.st_mode) else "f"
                })
        return sorted(dir_items, key=lambda item:(item["type"], item["name"]))

    def get_file_data(self, file_path):
        file_size = self._sftp_obj.stat(file_path).st_size
        with self._sftp_obj.open(file_path, 'rb') as f:
            f.prefetch(file_size)
            while True:
                data = f.read(1048576)

                if not data:
                    break

                yield data

    def put_file_data(self, local_file, remote_file):
        with self._sftp_obj.file(remote_file, "wb") as f:
            f.set_pipelined(True)
            while True:
                data = local_file.read(1048576)

                if not data:
                    break
                else:
                    f.write(data)


if __name__ == '__main__':
    sftp_obj = SFTP()
    sftp_obj.connect("192.168.30.80", "by", "by")
    # for item in sftp_obj.get_dir_items("/"):
    #     print(item["name"])

    # for data in sftp_obj.get_file_data("/home/by/a.txt"):
    #     print(data)

    with open("requirements.txt", "rb") as f:
        sftp_obj.put_file_data(f, "/home/by/requirements.txt")
